package com.kly.ai.app;

import cn.hutool.core.io.resource.ResourceUtil;
import com.kly.ai.advisor.MyLoggerAdvisor;
import com.kly.ai.advisor.ReReadingAdvisor;
import com.kly.ai.chatmemory.FileBasedChatMemory;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.List;

import static com.kly.ai.constants.FileConstant.MESSAGE_SAVE_DIR;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * @author admin
 * @date 2025/6/3
 */
@Component
@Slf4j
public class MyAgent {

    private final ChatClient myAgentClient;

    // @Resource
    // @Lazy
    // private ToolCallbackProvider toolCallbackProvider;

    @Resource
    private VectorStore psychologistVectorStore;


    @Resource
    private ToolCallback[] allTools;

    private final ChatClient.Builder chatClientBuilder;


    public MyAgent(ChatClient.Builder chatClientBuilder) {
        this.chatClientBuilder = chatClientBuilder;
        // 读取心理专家提示词
        String prompt = ResourceUtil.readUtf8Str("prompt/PsychologistPrompt.txt");
        FileBasedChatMemory fileBasedChatMemory = new FileBasedChatMemory(MESSAGE_SAVE_DIR);
        this.myAgentClient = chatClientBuilder
                .defaultSystem(prompt)
                .defaultAdvisors(new MessageChatMemoryAdvisor(fileBasedChatMemory))
                .build();
    }

    /**
     * AI对话
     *
     * @param userPrompt
     * @param conversationId
     * @return
     */
    public Flux<String> call(String conversationId, String userPrompt) {
        // 1.查询扩展
        MultiQueryExpander multiQueryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(chatClientBuilder)
                .includeOriginal(false)
                .numberOfQueries(3)
                .build();
        List<Query> queryList = multiQueryExpander.expand(new Query(userPrompt));
        // 2. 查询重写
        StringBuilder queryStringBuilder = new StringBuilder();
        for (Query query : queryList) {
            RewriteQueryTransformer rewriteQueryTransformer = RewriteQueryTransformer.builder()
                    .chatClientBuilder(chatClientBuilder)
                    .build();
            Query transformQuery = rewriteQueryTransformer.transform(query);
            queryStringBuilder.append(transformQuery).append("\n");
        }

        // 获取文档信息
        return myAgentClient
                .prompt()
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 20))
                .advisors(new ReReadingAdvisor())
                .advisors(new MyLoggerAdvisor())
                .advisors(new QuestionAnswerAdvisor(psychologistVectorStore))
                .tools(allTools)
                .user(userPrompt)
                .stream()
                .content();
    }
}
